package com.starhub.application.rag.store;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.starhub.application.model.biz.ScoringModelConnectionFactory;
import com.starhub.application.model.biz.VectorModelFactory;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.dto.RetrieveDto;
import com.starhub.application.rag.result.RagResult;
import com.starhub.application.rag.service.EmbeddingService;
import com.starhub.application.rag.service.MilvusConfigService;
import com.starhub.common.bean.model.MilvusConfig;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.extern.slf4j.Slf4j;

/**
 * 检索存储
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-25
 */
@Slf4j
@Component
public class RetrieveStore {

    @Autowired
    private MilvusConfigService milvusConfigService;

    @Autowired
    private LangchainStoreBuilder langchainStoreBuilder;

    @Autowired
    private ScoringModelConnectionFactory scoringModelFactory;

    @Autowired
    private EmbeddingService embeddingService;

    /**
     * 根据内容检索知识条目(内容检索模式)
     * 
     * @param retrieveDto
     * @return
     * @throws Exception 
     */
    public RagResult retrieve(RetrieveDto retrieveDto) throws Exception {
        // 进行向量化
        Embedding questionEmbedding = embeddingService.embedding(retrieveDto.getVectorModel(), retrieveDto.getQueryText());

        return retrieve(retrieveDto, questionEmbedding);
    }

    /**
     * 根据内容检索知识条目(内容检索模式)
     * 
     * @param retrieveDto
     * @param questionEmbedding
     * @return
     */
    public RagResult retrieve(RetrieveDto retrieveDto, Embedding questionEmbedding) {
        try {

            // 3. 获取Milvus配置和向量存储对象
            MilvusConfig config = milvusConfigService.getConfigByMark(RagConstant.MILVUS_MARK);
            EmbeddingStore<TextSegment> embeddingStore = langchainStoreBuilder.getOrCreateMilvusEmbeddingStore(config,
                    retrieveDto.getCollectionIdentifier(), retrieveDto.getDimension());

            // 4. 执行向量检索
            EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                    .queryEmbedding(questionEmbedding)
                    .maxResults(retrieveDto.getTopK())
                    .minScore(retrieveDto.getMinScore())
                    .build();

            EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);

            // 5. 对检索结果进行重排序
            if (retrieveDto.isNeedRerank()) {
                searchResult = rerank(retrieveDto.getQueryText(), searchResult, retrieveDto.getScoringModel());
            }

            RagResult ragResult = RagResult.builder()
                    .question(retrieveDto.getQueryText())
                    .searchResult(searchResult)
                    .build();
            return ragResult;
        } catch (Exception e) {
            throw new RuntimeException("向量检索失败: " + e.getMessage(), e);
        }
    }

    /**
     * 使用重排序模型对检索结果进行重新排序
     * 
     * @param question         用户问题
     * @param searchResults    原始检索结果
     * @param scoringModelMark 重排序模型标识
     * @return 重排序后的结果
     */
    private EmbeddingSearchResult<TextSegment> rerank(String question,
            EmbeddingSearchResult<TextSegment> searchResults,
            String scoringModelMark) {
        try {
            // 获取重排序模型实例
            ScoringModel scoringModel = scoringModelFactory.getScoringModel(scoringModelMark);// bcev1

            List<EmbeddingMatch<TextSegment>> matches = new ArrayList<>(searchResults.matches());
            List<TextSegment> segments = matches.stream()
                    .map(EmbeddingMatch::embedded)
                    .toList();

            // 使用重排序模型对文档进行评分
            List<Double> scores = scoringModel.scoreAll(segments, question).content();

            // 创建新的匹配列表，结合原始相似度和重排序分数
            Map<EmbeddingMatch<TextSegment>, Double> combinedScores = new HashMap<>();
            for (int i = 0; i < matches.size(); i++) {
                EmbeddingMatch<TextSegment> match = matches.get(i);
                double combinedScore = (scores.get(i) + match.score()) / 2.0;
                combinedScores.put(match, combinedScore);
            }

            // 根据新的分数重新排序
            matches.sort((a, b) -> Double.compare(
                    combinedScores.getOrDefault(b, b.score()),
                    combinedScores.getOrDefault(a, a.score())));

            // 创建新的搜索结果对象
            return new EmbeddingSearchResult<>(matches);
        } catch (Exception e) {
            log.error("重排序过程发生错误", e);
            // 发生错误时返回原始结果
            return searchResults;
        }
    }

}
